import json
from multiprocessing import cpu_count
from dataset_adapter.language_python import Py150kVarmisuseDatasetGenerator
from enum import Enum

class DataFold(Enum):
    TRAIN = 0
    VALIDATION = 1
    TEST = 2

def name_to_dataloader(name: str, path: str, data_fold: DataFold, vocab_dict,args, num_workers=int(cpu_count() / 2)):  # 返回类型为对象
    """将字符串转为数据预处理对象

    Args:
        name (str): 数据预处理的类名
        path (str): 数据目录地址
        args (_type_): 可能用到的参数

    Raises:
        ValueError: 类名不存在
    """
    name = name.lower()
    if name in ["csharp"]:
        datasetCls = CSharpStaticGraphDatasetGenerator
    elif name in ["python", "py150", "py"]:
        datasetCls = Py150kVarmisuseDatasetGenerator
    else:
        raise ValueError("Unkown dataset name '%s'" % name)

    # 返回数据库对象
    return datasetCls(
        path,
        vocab_dict,
        {},
        num_edge_types=args.num_edge_types,
        batch_size=args.batch_size,
        shuffle=True if data_fold == DataFold.TRAIN else False,
        graph_node_max_num_chars=args.graph_node_max_num_chars,
        max_graph=args.max_graph,
        max_variable_candidates=args.max_variable_candidates,
        max_node_per_graph=args.max_node_per_graph,
        num_workers=num_workers,
        device=args.device,
        slice_edge_type=args.slice_edge_type,
        slot_singal=args.slot_singal,
        task_type=args.task_type
    )